from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from PIL import Image
import numpy as np
import time
import sys
import urllib.parse

# 回调地址
redictURL = sys.argv[1]
# 请求链接
url = sys.argv[2]
# 1 长截图 2 窗口大小截图
type = sys.argv[3]
# 1 手机端 2 PC端
is_mobile = sys.argv[4] 
# 最大高度
scrollHeight = sys.argv[5]

def spider():
    driver_path = "/data/driver/chromedriver"
    options = Options()
    if is_mobile == 1:
        # 调整为手机端
        mobile_emulation = {"deviceName": "iPhone 12 Pro"}
        options.add_experimental_option("mobileEmulation", mobile_emulation)
        options.add_experimental_option("excludeSwitches", ["enable-logging"])
    options.add_argument('headless')
    driver = webdriver.Chrome(executable_path=driver_path,options=options)
    try:
        driver.get(url)
        driver.maximize_window()
        # time.sleep(3)
        # width = driver.execute_script("return document.documentElement.scrollWidth")
        window_height = driver.get_window_size()['height']
        page_height = driver.execute_script("return document.documentElement.scrollHeight")
        time.sleep(1)
        driver.save_screenshot('qq.png')
        if type == 2:
            # 直接返回截到的第一张图即可
            return sys.path[0]+"qq.png"
        if page_height > window_height:
            n = page_height // window_height  # 需要滚动的次数
            base_mat = np.atleast_2d(Image.open('qq.png'))  # 打开截图并转为二维矩阵
            for j in range(n):
                driver.execute_script(f'document.documentElement.scrollTop={window_height*(j+1)};')
                time.sleep(.5)
                driver.save_screenshot(f'qq_{j}.png')  # 保存截图
                mat = np.atleast_2d(Image.open(f'qq_{j}.png'))  # 打开截图并转为二维矩阵
                base_mat = np.append(base_mat, mat, axis=0)  # 拼接图片的二维矩阵
            Image.fromarray(base_mat).save('./Image/image.png')
            driver.close()
            driver.quit()
            return sys.path[0]+"Image/image.png"
            # send_file("Image/image.png",sys.path[0])
        return sys.path[0]+"qq.png"
        # send_file("qq.png",sys.path[0])
    except Exception as e:
        driver.close()
        driver.quit()

if __name__ == '__main__':
    spider()

